import logging
import os
import wave
from typing import IO, Union

import jpype

import core.model_runtime.model_providers.picc.helper.jar_loader as JarLoader
from configs import dify_config

piccJarPath = dify_config.PICC_JAR_PATH
piccClientJarPath = dify_config.PICC_CLIENT_JAR_PATH
JarLoader.loadJar(jarPath=piccJarPath)
JarLoader.loadJar(jarPath=piccClientJarPath)
appCredentialClass = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.common.entity.AppCredential")
iatApiclientClass = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.asr.client.IatApiClient")
apiRequestClass = JarLoader.getClass("com.iflytek.ai.platform.api.core.entity.ApiRequest")
apiResponseClass = JarLoader.getClass("com.iflytek.ai.platform.api.core.entity.ApiResponse")
onlineIatRequestClass = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.asr.entity.OnlineIatRequest")
onlineIatResponseClass = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.asr.entity.OnlineIatResponse")
ttsRequestRequestClass = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.tts.entity.TtsRequest")
onlineIatListenerClass = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.asr.listener.OnlineIatListener")
byteArrayInputStreamClass = JarLoader.getClass("java.io.ByteArrayInputStream")
fileInputStreamClass = JarLoader.getClass("java.io.FileInputStream")
asrClientClass = JarLoader.getClass("cn.com.picc.asr.ASRClient")
ttsClientClass = JarLoader.getClass("cn.com.picc.tts.TTSClient")

ttsRequest = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.tts.entity.TtsRequest")
ttsClient = JarLoader.getClass("com.iflytek.ai.platform.api.sdk.tts.client.TtsApiClient")
# void = None
logger = logging.getLogger(__name__)


def pcm_to_wav(pcm_file, wav_file, channels=1, sample_width=2, sample_rate=44100):
    pcm_data = open(pcm_file, 'rb').read()

    wav_file = wave.open(wav_file, 'wb')
    wav_file.setnchannels(channels)
    wav_file.setsampwidth(sample_width)
    wav_file.setframerate(sample_rate)
    wav_file.writeframes(pcm_data)
    wav_file.close()


class ASRClient:

    client = None

    timeout = 10

    def __init__(self, url: str,
                 ability_id: str,
                 app_id: str,
                 secret_id: str,
                 secret_key: str,
                 timeout: int = 10):
        AbilityId = jpype.JString(ability_id)
        AppId = jpype.JString(app_id)
        SecretId = jpype.JString(secret_id)
        SecretKey = jpype.JString(secret_key)
        self.client = asrClientClass(url, AppId, AbilityId, SecretId, SecretKey)
        self.timeout = timeout

    def invoke(self,
               bytes: IO[bytes] = None,
               filePath: str = None,
               rate: int = 16000,
               bits: int = 10) -> str:
        input = ""
        if bytes is not None:
            byteArray = bytes.read()
            input = jpype.JArray(jpype.JByte)(byteArray)
        else:
            input = jpype.JString(filePath)
        result = self.client.invoke(input, jpype.JInt(rate), jpype.JInt(bits), jpype.JInt(self.timeout))
        return str(result)


class TTSClient:

    client = None

    def __init__(self, url: str,
                 ability_id: str,
                 app_id: str,
                 secret_id: str,
                 secret_key: str,):
        AbilityId = jpype.JString(ability_id)
        AppId = jpype.JString(app_id)
        SecretId = jpype.JString(secret_id)
        SecretKey = jpype.JString(secret_key)
        self.client = ttsClientClass(url, AppId, AbilityId, SecretId, SecretKey)

    def invoke(self,
               text: str,
               to_file_path: str = None,
               voiceId: str = "67910",
               volume: int = 10,
               speed: int = 0,
               pitch: int = 0
               ) -> Union[bytes, bool]:
        text = jpype.JString(text)
        voiceId = jpype.JString(voiceId)

        if to_file_path is None:
            # 收到的是java的 byte[]
            result = self.client.invoke(text, voiceId, jpype.JInt(volume), jpype.JInt(speed), jpype.JInt(pitch))
            res = bytes(i % 256 for i in result)
            return res

        else:
            to_file_path_pcm = to_file_path.split(".")[0] + ".pcm"
            return_file = self.client.invoke(text, voiceId, jpype.JString(to_file_path_pcm), jpype.JInt(volume), jpype.JInt(speed), jpype.JInt(pitch))
            if return_file:
                pcm_to_wav(to_file_path_pcm, to_file_path)
                os.remove(to_file_path_pcm)
            return return_file
